from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import matplotlib.pyplot as plt
import numpy as np


class GRU_test(nn.Module):
    out_dims: int = 64
    hidden_dims: int = 64

    @nn.compact
    def __call__(self, state, input):
        new_state, output = nn.GRUCell()(state, input)
        return new_state, output

    def initial_state(self, batch_size):
        # Zero initialization
        return nn.GRUCell.initialize_carry(jax.random.PRNGKey(0), (batch_size, ), self.hidden_dims,
                                            init_fn=nn.initializers.zeros)

    def initial_state_rnd(self, batch_size, key):
        # random initialization
        state = jax.random.normal(key, (batch_size, self.hidden_dims))
        return state

    @staticmethod
    def state_metrics(state):
        return {}

@partial(jax.jit, static_argnums=(2,))
def init_fn(seed, input_, model):
    """ init function for single model
    """
    params = model.init(seed, model.initial_state(input_.shape[0]), input_)
    return params

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def transform(X, W, B):
    h = jnp.dot(X, W) + B
    return jnp.tanh(h)

# compute cosine similarity of two vectors
@partial(jax.jit)
def cos_sim(x, y):
    return jnp.dot(x, y) / (jnp.linalg.norm(x) * jnp.linalg.norm(y))

@partial(jax.jit)
def KL_divergence(p, q):
    return jnp.sum(p * jnp.log(p / q))

def main():

    nn_size = 20

    len_trail1 = 100

    mean_influences = []

    n_tests = 50

    # gru_network = GRU_test(hidden_dims=nn_size, out_dims=nn_size)
    # # make param of the network
    # input_ = jnp.ones((1, nn_size))
    # params = init_fn(jax.random.PRNGKey(np.random.randint(0, 100000)), input_, gru_network)
    # state_ = gru_network.initial_state(input_.shape[0])
    # state_, y1 = model_forward(params, state_, input_, gru_network)
    # print("shape of y1 = ", y1.shape)

    for k in range(n_tests):

        # len_trail1 += 100
        influence = []

        for t in range(100):
            
            '''
            generate the input sequence for the first trail
            '''
            input_seq = []
            for i in range(len_trail1):
                input_seq.append(jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size)))

            '''
            generate params of the neural network
            '''
            # create a random vector of size 20 with values between 0 and 1
            X = jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size))
            X0 = X.copy()
            # create a random matrix of size 20x20 with values between 0 and 1
            W = jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (nn_size, nn_size))
            # create a random vector of size 20 with values between 0 and 1
            B = jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size))

            # gru_network = GRU_test(hidden_dims=nn_size, out_dims=nn_size)
            # params = init_fn(jax.random.PRNGKey(np.random.randint(0, 100000)), input_seq[0], gru_network)
            # state_ = gru_network.initial_state(input_seq[i].shape[0])
            
            '''
            compute the final state of the neural network for the first trail
            '''
            for i in range(len_trail1):
                X = transform(X + input_seq[i], W, B)

            # for i in range(len_trail1):
            #     state_, output_ = model_forward(params, state_, input_seq[i], gru_network)
        
            state1 = X.copy()

            # state1 = output_.copy()

            '''
            make a small perturbation to the initial state for the second trail
            '''
            X = X0.copy()
            # create a purturbation vector of size 20 with values between 0 and 1
            X_purt = jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size))
            X = X + X_purt

            # state_ = gru_network.initial_state(input_seq[i].shape[0])
            # input_purt = jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size))
            # input_seq[0] = input_seq[0] + input_purt
            # state_ = state_ + input_purt

            for n in range(input_seq[i].shape[0]):
                input_purt = jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size))
                input_seq[n] = input_seq[n] + input_purt

            # input_seq1 = []
            # for n in range(len_trail1):
            #     input_seq1.append(input_seq[n] + jax.random.normal(jax.random.PRNGKey(np.random.randint(0, 100000)), (1, nn_size)))

            '''
            compute the final state of the neural network for the second trail
            '''
            for i in range(len_trail1):
                X = transform(X + input_seq[i], W, B)

            # for i in range(len_trail1):
            #     state_, output_ = model_forward(params, state_, input_seq[i], gru_network)

            state2 = X.copy()

            # state2 = output_.copy()

            '''
            compute the cosine similarity between the final states of the two trails
            '''
            state1 = state1.reshape(-1)
            state2 = state2.reshape(-1)
            similarity = cos_sim(state1, state2)

            # print(jnp.mean(X_purt), similarity)
            print(similarity)

            influence.append(similarity)

        nn_size += 20

        mean_influences.append(np.mean(influence))
        print("len_trail = ", len_trail1)
        print("nn_size = ", nn_size)

    _, ax = plt.subplots()
    
    plt.plot(mean_influences)
    
    labels = [20 * i for i in range(n_tests)]
    ax.set_xticklabels(labels)
    
    # name the x axis "trial length"
    plt.xlabel('nn_size')
    plt.title('The effect of initial random perturbation on the final state versus nn_size')
    plt.show()


if __name__ == "__main__":
    main()